import os
import random
import json
from concurrent.futures import ThreadPoolExecutor

def generate_replace_rule():
    """
    生成固定长度的字节对字节一对一替换规则
    """
    original_bytes = list(range(256))
    random.shuffle(original_bytes)
    replacement_bytes = list(range(256))
    random.shuffle(replacement_bytes)
    replace_rule = {}
    for i in range(256):
        # 确保替换规则一对一，不重复使用原始字节和替换字节
        replace_rule[original_bytes[i]] = replacement_bytes[i]
    return replace_rule

def replace_header_content(data, replace_rule):
    """
    根据给定的替换规则对文件头部二进制内容进行替换，保持数据长度不变
    """
    new_header = bytearray(data)
    for index in range(len(data)):
        byte_value = new_header[index]
        if byte_value in replace_rule:
            new_header[index] = replace_rule[byte_value]
    return bytes(new_header)

def process_file(file_path, replace_rule, replace_length):
    try:
        file_size = os.path.getsize(file_path)
        with open(file_path, 'r+b') as file:
            # 只读取文件头部指定长度的数据（这里是256字节）
            header_data = file.read(replace_length)
            replaced_header = replace_header_content(header_data, replace_rule)
            # 将文件指针移动回文件开头
            file.seek(0)
            # 写回替换后的头部数据
            file.write(replaced_header)
            # 如果文件长度大于256字节，需要将文件指针移动到原头部末尾（256字节处）
            if file_size > replace_length:
                file.seek(replace_length)
    except Exception as e:
        print(f"处理文件 {file_path} 时出错: {e}")

def encrypt_files(folder_path, replace_length=256, max_workers=12):
    """
    对文件夹下所有文件头部前指定长度的二进制内容进行随机替换，并保存全局反向替换规则到JSON文件
    """
    # 生成替换规则
    replace_rule = generate_replace_rule()
    # 生成反向替换规则，方便解密时直接使用
    reverse_replace_rule = {v: k for k, v in replace_rule.items()}
    # 生成替换规则保存的文件名，与文件夹路径相关
    rule_file_name = os.path.join(os.path.dirname(folder_path), "replace_rule.json")
    # 将反向替换规则保存到JSON文件
    with open(rule_file_name, 'w+') as rule_file:
        json.dump(reverse_replace_rule, rule_file)

    file_paths = []
    for root, dirs, files in os.walk(folder_path):
        for file_name in files:
            file_path = os.path.join(root, file_name)
            file_paths.append(file_path)

    # 使用线程池处理文件
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [executor.submit(process_file, file_path, replace_rule, replace_length) for file_path in file_paths]
        for future in futures:
            future.result()

    return rule_file_name

if __name__ == "__main__":
    folder_to_encrypt = "testdir"  # 替换为实际要加密的文件夹路径
    rule_file_path = encrypt_files(folder_to_encrypt, 4194304)  # 使用自定义分块大小进行加密
    print(f"替换规则已保存到 {rule_file_path}")